import linecache

import PIL.Image as Image
import torch
from torch.utils.data import Dataset
from itertools import chain
import numpy as np
from myutils.mytransformers import clip_1_99,color_0_255,img_pad_to_square


from torchvision import transforms
import skimage
import cv2

from myutils.IndexGetter import IndexGetter

# 单条数据加载
class Dataset_one(Dataset):
    """单模型数据加载

    Input:
        path_txt    :txt目录文件路径
        label_count :每一种label的数量,[1,2,3,4...]
        transform   :图片转换
    Else:
        label_presum    :前缀和（各label的起始索引）,[0,1,3,6...]
        index_list      :对每个label，都有1个索引列表,[3,2,4,0,1]*24
        index_now       :对每个label，都有索引数组
        label_count_sum :数据行数
        length          :max(label_count)*len(label_count)
        label_location  :label 转 location
    
    """
    
    def __init__(self,path_txt,label_count,is_test,use_shape,use_location):
        self.path_txt = path_txt # 目录路径
        self.is_test      = is_test
        self.use_shape    = use_shape
        self.use_location = use_location

        self.index_list = list(chain(*map(lambda i:[i[0]]*i[1],enumerate(label_count))))
        
        if not is_test :
            self.indexGetter = IndexGetter(list(enumerate(self.index_list)),24)

        # laebl -> location: [0,1,2]->[中,亚中,端(带随提近端)]
        #                                      [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,'X','Y']
        self.label_location = torch.LongTensor([0,0,0,1,1,1,1,1,1, 1, 1, 1, 2, 2, 2, 0, 1, 1, 0, 0, 2, 2,  1,  2])

        
        # 图片转换
        transforms_ext = [] if  is_test else [
            # transforms.RandomHorizontalFlip(),
            # transforms.RandomVerticalFlip(),
            transforms.RandomRotation((0,360),expand=False,fill=255),]  # 训练时的随机变化，测试时不需要
        self.transform=transforms.Compose([
            transforms.ToTensor(),
            clip_1_99,
            color_0_255,
            img_pad_to_square,
            transforms.Resize((64,64)),
            *transforms_ext,
            ])
    def __len__(self):
        return len(self.index_list) if self.is_test else len(self.indexGetter)
    def __getitem__(self,index): 
        """获得指定记录

        Input:
            index: 第几条数据，这里只是表示一个连续增长的数，用于判断 `index%24`
        Output:
            img      :图片
            shape    :大小
            Y        :label标签
            location :着丝粒位置
        Else:
            label_target   :目标图片的label
            index_of_list  :当前label索引的 下一个下标
            index_of_label :取出label索引
            index_of_txt   :在文本中的索引
        """
        if not self.is_test :
            index = self.indexGetter.get_index(index%self.indexGetter.class_number)

        (path_img,Y,location) = self.__get_path__(index)
        (img,shape) = self.__load_by_path__(path_img)

        if not self.use_shape:
            shape = torch.tensor([0,0,0,0])
        if not self.use_location:
            location = 0

        return (img,shape,Y,location)
    
    # 
    def __get_path__(self,index):
        """获得指定记录的路径，类型，着丝粒位置

        Input:
            index: 第几条数据
        Output:
            path_img :图片路径
            Y        :label标签
            location :着丝粒位置
        Else:
            line:一条记录（一行就是一条记录，path_img + label）
        """
        line = linecache.getline(self.path_txt,index+1)
        # print(index,line)
        path_img,Y = line.split(' ')
        Y  = torch.LongTensor([int(Y)]) # shape=(1)
        location = self.label_location[Y].reshape(1) # shape=(1)
        return (path_img,Y,location)
    
    # （shape，transforms）
    def __load_by_path__(self,path_img):
        """通过路径加载图片"""
        img_PIL  = Image.open(path_img)
        W,H = img_PIL.size[-2:]

        img  = self.transform(img) # shape=(3,224,224)

        len = self.skeletonize(img).sum()
        shape  = torch.tensor((len/64*H,len,W,H),dtype=torch.float) # shape=(2)

        return (img,shape)
    

    
    def skeletonize(self,img):
        kernel=np.ones((3,3))
        img1_PIL = transforms.ToPILImage()(img)
        img1_01 = np.array(img1_PIL)>14
        img1_01 = img1_01.astype(np.uint8)
        for i in range(1):
            img1_01=cv2.morphologyEx(img1_01,cv2.MORPH_CLOSE,kernel=kernel)
        img3 = skimage.morphology.skeletonize(np.where(img1_01,1,0)) #*255
        return img3
    

